"""
模型对比API路由
提供多模型对比相关的RESTful接口和WebSocket端点
"""

from fastapi import APIRouter, HTTPException, Depends, WebSocket, WebSocketDisconnect, Query
from fastapi.responses import JSONResponse
from typing import List, Optional, Dict, Any
from datetime import datetime
import asyncio
import json
import uuid
import logging

from motor.motor_asyncio import AsyncIOMotorDatabase

from ..auth.jwt_handler import get_current_user
from ..core.database import get_database
from ..models.comparison import (
    ComparisonRequest,
    ComparisonResponse,
    ComparisonSession,
    ComparisonStatus,
    ModelResult,
    ModelSelectionEvent,
    ComparisonHistory
)
from ..repositories.comparison_repository import ComparisonRepository
from ..comparisons.parallel_generator import (
    ParallelModelGenerator,
    ModelConfig,
    ModelStatus
)
from ..scoring.comparison_scorer import ComparisonScorer

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/comparisons", tags=["comparisons"])

# WebSocket连接管理器
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, List[WebSocket]] = {}
        
    async def connect(self, websocket: WebSocket, session_id: str):
        await websocket.accept()
        if session_id not in self.active_connections:
            self.active_connections[session_id] = []
        self.active_connections[session_id].append(websocket)
        logger.info(f"WebSocket连接建立: {session_id}")
        
    def disconnect(self, websocket: WebSocket, session_id: str):
        if session_id in self.active_connections:
            self.active_connections[session_id].remove(websocket)
            if not self.active_connections[session_id]:
                del self.active_connections[session_id]
        logger.info(f"WebSocket连接断开: {session_id}")
        
    async def send_message(self, message: str, session_id: str):
        if session_id in self.active_connections:
            for connection in self.active_connections[session_id]:
                try:
                    await connection.send_text(message)
                except:
                    pass
                    
    async def broadcast(self, message: str, session_id: str):
        await self.send_message(message, session_id)

manager = ConnectionManager()

@router.post("/create", response_model=ComparisonResponse)
async def create_comparison(
    request: ComparisonRequest,
    current_user: dict = Depends(get_current_user),
    db: AsyncIOMotorDatabase = Depends(get_database)
):
    """
    创建新的模型对比任务
    
    Args:
        request: 对比请求
        current_user: 当前用户
        db: 数据库连接
        
    Returns:
        对比响应
    """
    try:
        # 创建会话ID
        session_id = f"comp_{datetime.utcnow().strftime('%Y%m%d')}_{uuid.uuid4().hex[:8]}"
        
        # 创建对比会话
        session = ComparisonSession(
            session_id=session_id,
            user_id=current_user["user_id"],
            prompt=request.prompt,
            system_prompt=request.system_prompt,
            models=request.models,
            metadata=request.metadata,
            status=ComparisonStatus.PROCESSING
        )
        
        # 保存到数据库
        repo = ComparisonRepository(db)
        await repo.create_session(session)
        
        # 启动异步生成任务
        asyncio.create_task(
            process_comparison(session_id, request, db)
        )
        
        return ComparisonResponse(
            session_id=session_id,
            status=ComparisonStatus.PROCESSING,
            results=[],
            aggregation={"message": "对比任务已创建，正在处理中"}
        )
        
    except Exception as e:
        logger.error(f"创建对比任务失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

async def process_comparison(
    session_id: str,
    request: ComparisonRequest,
    db: AsyncIOMotorDatabase
):
    """
    异步处理对比任务
    
    Args:
        session_id: 会话ID
        request: 对比请求
        db: 数据库连接
    """
    try:
        repo = ComparisonRepository(db)
        generator = ParallelModelGenerator()
        
        # 准备模型配置
        model_configs = [
            ModelConfig(
                name=model.name,
                version=model.version,
                parameters=model.parameters,
                timeout=30.0
            )
            for model in request.models
        ]
        
        # 进度回调
        async def progress_callback(model_name: str, status: ModelStatus):
            message = {
                "type": "progress",
                "model": model_name,
                "status": status.value,
                "timestamp": datetime.utcnow().isoformat()
            }
            await manager.broadcast(json.dumps(message), session_id)
        
        # 执行并行生成
        results = await generator.generate_parallel(
            prompt=request.prompt,
            model_configs=model_configs,
            system_prompt=request.system_prompt,
            callback=progress_callback
        )
        
        # 评分
        scorer = ComparisonScorer()
        
        # 保存结果
        for result in results:
            # 计算评分
            scores = await scorer.score_content(
                content=result.content,
                prompt=request.prompt
            ) if result.status == ModelStatus.COMPLETED else {}
            
            model_result = ModelResult(
                model=result.model,
                content=result.content,
                tokens=result.tokens,
                time=result.time_taken,
                score=scores,
                status=result.status.value,
                error=result.error
            )
            
            await repo.add_result(session_id, model_result)
            
            # 发送结果更新
            message = {
                "type": "result",
                "model": result.model,
                "result": model_result.dict(),
                "timestamp": datetime.utcnow().isoformat()
            }
            await manager.broadcast(json.dumps(message), session_id)
        
        # 聚合结果
        aggregation = await generator.aggregate_results(results)
        
        # 更新会话状态
        all_success = all(r.status == ModelStatus.COMPLETED for r in results)
        final_status = ComparisonStatus.COMPLETED if all_success else ComparisonStatus.PARTIAL
        await repo.update_status(session_id, final_status)
        
        # 发送完成消息
        message = {
            "type": "complete",
            "status": final_status.value,
            "aggregation": aggregation,
            "timestamp": datetime.utcnow().isoformat()
        }
        await manager.broadcast(json.dumps(message), session_id)
        
    except Exception as e:
        logger.error(f"处理对比任务失败 {session_id}: {e}")
        
        # 更新为失败状态
        repo = ComparisonRepository(db)
        await repo.update_status(session_id, ComparisonStatus.FAILED)
        
        # 发送错误消息
        message = {
            "type": "error",
            "error": str(e),
            "timestamp": datetime.utcnow().isoformat()
        }
        await manager.broadcast(json.dumps(message), session_id)

@router.get("/{comparison_id}", response_model=ComparisonSession)
async def get_comparison(
    comparison_id: str,
    current_user: dict = Depends(get_current_user),
    db: AsyncIOMotorDatabase = Depends(get_database)
):
    """
    获取对比结果
    
    Args:
        comparison_id: 对比会话ID
        current_user: 当前用户
        db: 数据库连接
        
    Returns:
        对比会话
    """
    try:
        repo = ComparisonRepository(db)
        session = await repo.get_session(comparison_id)
        
        if not session:
            raise HTTPException(status_code=404, detail="对比会话不存在")
            
        # 验证权限
        if session.user_id != current_user["user_id"]:
            raise HTTPException(status_code=403, detail="无权访问此对比")
            
        return session
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取对比结果失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/{comparison_id}/select")
async def select_model(
    comparison_id: str,
    model: str,
    reason: Optional[str] = None,
    current_user: dict = Depends(get_current_user),
    db: AsyncIOMotorDatabase = Depends(get_database)
):
    """
    选择最优模型
    
    Args:
        comparison_id: 对比会话ID
        model: 选择的模型名称
        reason: 选择原因
        current_user: 当前用户
        db: 数据库连接
        
    Returns:
        操作结果
    """
    try:
        repo = ComparisonRepository(db)
        
        # 验证会话存在和权限
        session = await repo.get_session(comparison_id)
        if not session:
            raise HTTPException(status_code=404, detail="对比会话不存在")
        if session.user_id != current_user["user_id"]:
            raise HTTPException(status_code=403, detail="无权操作")
            
        # 更新选择
        await repo.set_selected_model(comparison_id, model)
        
        # 记录选择事件（A/B测试）
        event = ModelSelectionEvent(
            session_id=comparison_id,
            user_id=current_user["user_id"],
            selected_model=model,
            reason=reason,
            scores={r.model: r.score.get("overall", 0) for r in session.results if r.score}
        )
        await repo.record_selection_event(event)
        
        return {"status": "success", "message": f"已选择 {model}"}
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"选择模型失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/history", response_model=ComparisonHistory)
async def get_history(
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
    status: Optional[ComparisonStatus] = None,
    current_user: dict = Depends(get_current_user),
    db: AsyncIOMotorDatabase = Depends(get_database)
):
    """
    获取对比历史记录
    
    Args:
        page: 页码
        page_size: 每页大小
        status: 状态过滤
        current_user: 当前用户
        db: 数据库连接
        
    Returns:
        对比历史
    """
    try:
        repo = ComparisonRepository(db)
        
        offset = (page - 1) * page_size
        sessions = await repo.get_user_history(
            user_id=current_user["user_id"],
            limit=page_size,
            offset=offset,
            status_filter=status
        )
        
        # 获取总数
        total_count = await db.model_comparisons.count_documents(
            {"user_id": current_user["user_id"]}
        )
        
        return ComparisonHistory(
            user_id=current_user["user_id"],
            sessions=sessions,
            total_count=total_count,
            page=page,
            page_size=page_size
        )
        
    except Exception as e:
        logger.error(f"获取历史记录失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@router.websocket("/ws/{comparison_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    comparison_id: str
):
    """
    WebSocket端点，用于实时推送对比进度
    
    Args:
        websocket: WebSocket连接
        comparison_id: 对比会话ID
    """
    await manager.connect(websocket, comparison_id)
    
    try:
        # 发送连接成功消息
        await websocket.send_text(json.dumps({
            "type": "connected",
            "session_id": comparison_id,
            "timestamp": datetime.utcnow().isoformat()
        }))
        
        # 保持连接
        while True:
            data = await websocket.receive_text()
            
            # 处理ping消息
            if data == "ping":
                await websocket.send_text("pong")
                
    except WebSocketDisconnect:
        manager.disconnect(websocket, comparison_id)
        logger.info(f"WebSocket断开: {comparison_id}")
    except Exception as e:
        logger.error(f"WebSocket错误: {e}")
        manager.disconnect(websocket, comparison_id)

@router.get("/statistics/ab-testing")
async def get_ab_testing_stats(
    days: int = Query(7, ge=1, le=90),
    current_user: dict = Depends(get_current_user),
    db: AsyncIOMotorDatabase = Depends(get_database)
):
    """
    获取A/B测试统计数据
    
    Args:
        days: 统计天数
        current_user: 当前用户
        db: 数据库连接
        
    Returns:
        统计数据
    """
    try:
        repo = ComparisonRepository(db)
        stats = await repo.get_selection_statistics(
            user_id=current_user["user_id"],
            days=days
        )
        
        return stats
        
    except Exception as e:
        logger.error(f"获取统计数据失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))